#pragma once

/*
 * Optimized Skip List implementation
 */

#include <iostream>
#include <vector>
#include <random>
#include <limits>
#include <memory>
#include <algorithm>

template<typename K, typename V>
class SkipList {
private:
    struct Node {
        K key;
        V value;
        std::vector<Node*> forward;
        
        Node(K k, V v, int level) : key(std::move(k)), value(std::move(v)), forward(level, nullptr) {}
    };

    Node* head;
    const int max_level;
    int current_level;
    const float p;
    std::mt19937 gen;
    std::uniform_real_distribution<float> dis;

    int random_level() {
        int lvl = 1;
        while (dis(gen) < p && lvl < max_level) {
            lvl++;
        }
        return lvl;
    }

public:
    SkipList(int max_lvl = 32, float prob = 0.5) 
        : max_level(max_lvl), 
          current_level(1), 
          p(prob),
          gen(std::random_device{}()),
          dis(0.0f, 1.0f) {
        K min_key = std::numeric_limits<K>::min();
        head = new Node(min_key, V(), max_level);
    }

    ~SkipList() {
        Node* current = head;
        while (current) {
            Node* next = current->forward[0];
            delete current;
            current = next;
        }
    }

    // Disable copy operations
    SkipList(const SkipList&) = delete;
    SkipList& operator=(const SkipList&) = delete;

    void insert(K key, V value) {
        std::vector<Node*> update(max_level, nullptr);
        Node* current = head;

        // Search for the insertion point
        for (int i = current_level - 1; i >= 0; --i) {
            while (current->forward[i] && current->forward[i]->key < key) {
                current = current->forward[i];
            }
            update[i] = current;
        }

        current = current->forward[0];

        // If key exists, update value
        if (current && current->key == key) {
            current->value = std::move(value);
            return;
        }

        // Create new node
        int new_level = random_level();
        Node* new_node = new Node(std::move(key), std::move(value), new_level);

        // Update pointers
        for (int i = 0; i < new_level; ++i) {
            if (i >= current_level) {
                update[i] = head;
            }
            new_node->forward[i] = update[i]->forward[i];
            update[i]->forward[i] = new_node;
        }

        current_level = std::max(current_level, new_level);
    }

    bool search(const K& key, V& value) const {
        Node* current = head;

        for (int i = current_level - 1; i >= 0; --i) {
            while (current->forward[i] && current->forward[i]->key < key) {
                current = current->forward[i];
            }
        }

        current = current->forward[0];

        if (current && current->key == key) {
            value = current->value;
            return true;
        }

        return false;
    }

    bool contains(const K& key) const {
        Node* current = head;

        for (int i = current_level - 1; i >= 0; --i) {
            while (current->forward[i] && current->forward[i]->key < key) {
                current = current->forward[i];
            }
        }

        current = current->forward[0];
        return current && current->key == key;
    }

    bool remove(const K& key) {
        std::vector<Node*> update(max_level, nullptr);
        Node* current = head;

        for (int i = current_level - 1; i >= 0; --i) {
            while (current->forward[i] && current->forward[i]->key < key) {
                current = current->forward[i];
            }
            update[i] = current;
        }

        current = current->forward[0];

        if (!current || current->key != key) {
            return false;
        }

        // Update forward pointers
        for (int i = 0; i < current_level; ++i) {
            if (update[i]->forward[i] != current) {
                break;
            }
            update[i]->forward[i] = current->forward[i];
        }

        delete current;

        // Update current level
        while (current_level > 1 && head->forward[current_level - 1] == nullptr) {
            --current_level;
        }

        return true;
    }

    void display() const {
        for (int i = current_level - 1; i >= 0; --i) {
            Node* node = head->forward[i];
            std::cout << "Level " << i << ": ";
            while (node != nullptr) {
                std::cout << node->key << ":" << node->value << " ";
                node = node->forward[i];
            }
            std::cout << "\n";
        }
    }

    bool empty() const {
        return head->forward[0] == nullptr;
    }
};